# -*- coding: utf-8 -*-
#  @author  Bink
#  @date  2021/3/30 10:32 上午
#  @Email : 2641032316@qq.com


import os
import random
import numpy as np

import cv2

# images_path = '17flowers'
Angles = [5, 10, 15, 20, 60, 90, 100, 180]
Trans_Select_images = 10


def get_files(images_path):
    file_list = []
    for root, dirs, files in os.walk(images_path):
        # for dir in dirs:
        # print(os.path.join(root,dir))
        for file in files:
            if file.endswith('.jpg'):
                file_list.append(os.path.join(root, file))
    return file_list


def read_images(images_path):
    images_name = get_files(images_path)
    images = []
    for image_path in images_name:
        image = cv2.imread(image_path)
        images.append(image)
    # print('images: ', images)
    # print('images_name', images_name)
    return images, images_name


def resize_image(image):
    image_inter_area = cv2.resize(image, (128, 128), interpolation=cv2.INTER_AREA)
    image_inter_cubic = cv2.resize(image, (256, 256), interpolation=cv2.INTER_CUBIC)
    image_inter_nearest = cv2.resize(image, (300, 300), interpolation=cv2.INTER_NEAREST)
    return image_inter_area, image_inter_cubic, image_inter_nearest


def get_resize_images(images, images_name, resize_image_path):
    i = 0
    for image in images:
        # image = cv2.imread('data/111.jpg', 0)
        image_inter_area, image_inter_cubic, image_inter_nearest = resize_image(image)
        # print((resize_image_path + 'resize_inter_area_' + images_name[i].split('.')[0] + '.jpg').replace('/',"_")).replace('/',"_"))
        cv2.imwrite(resize_image_path + ('resize_inter_area_' + images_name[i].split('/')[-1]).replace('/', "_"),
                    image_inter_area)
        cv2.imwrite(resize_image_path + ('resize_inter_cubic_' + images_name[i].split('/')[-1]).replace('/', "_"),
                    image_inter_cubic)
        cv2.imwrite(resize_image_path + ('resize_inter_nearest_' + images_name[i].split('/')[-1]).replace('/', "_"),
                    image_inter_nearest)
        i = i + 1


def rotate_image(image, angle):
    (height, width) = image.shape[:2]
    center = (height // 2, width // 2)
    matrix = cv2.getRotationMatrix2D(center, angle, 1)
    # 旋转图像
    rotate_image = cv2.warpAffine(image, matrix, (width, height))
    return rotate_image


def get_rotate_images(images, images_name, rotate_image_path):
    i = 0
    for idx_image in range(len(images)):

        for angle in Angles:
            r_image = rotate_image(images[idx_image], angle)
            rotate_image_name = images_name[i].split('/')[-1].split('.')[0] + str('_rot_') + str(angle) + '.jpg'
            cv2.imwrite(rotate_image_path + rotate_image_name, r_image)
        i = i + 1


def translate_image(image, x_shift, y_shift):
    (height, width) = image.shape[:2]
    # 平移矩阵(浮点数类型)  x_shift +右移 -左移  y_shift -上移 +下移
    matrix = np.float32([[1, 0, x_shift], [0, 1, y_shift]])
    # 平移图像
    trans_image = cv2.warpAffine(image, matrix, (width, height))
    return trans_image


def get_trans_images(images, images_name, trans_image_path):
    np.random.seed(2)
    i = 0
    for image in images:
        # 获得随机平移坐标
        x_shift = np.random.randint(-8, 8, 1)
        y_shift = np.random.randint(-8, 8, 1)
        trans_image = translate_image(images[i], x_shift, y_shift)
        # 保存平移图片和平移坐标
        trans_image_name = images_name[i].split('/')[-1].split(".")[0] + '_shiftx' + str(x_shift[0]) + '_shifty' + str(
            y_shift[0]) + '.jpg'
        cv2.imwrite(trans_image_path + trans_image_name, trans_image)
        i = i + 1


def GaussianNoise(src, means, sigma, percetage):
    Noiseimage = src
    NoiseNum = int(percetage * src.shape[0] * src.shape[1])
    for i in range(NoiseNum):
        randX = random.randint(0, src.shape[0] - 1)
        randY = random.randint(0, src.shape[1] - 1)
        Noiseimage[randX, randY] = Noiseimage[randX, randY] + random.gauss(means, sigma)
        if Noiseimage[randX, randY].any() < 0:
            Noiseimage[randX, randY] = 0
        elif Noiseimage[randX, randY].all() > 255:
            Noiseimage[randX, randY] = 255
    return Noiseimage


def get_GaussianNoise_images(images, images_name, gaosi_image_path):
    i = 0
    for image in images:
        image = GaussianNoise(image, 2, 4, 0.8)
        cv2.imwrite(gaosi_image_path + 'gaussian_' + images_name[i].split('/')[-1], image)
        i = i + 1


def get_gray_images(images, images_name, gray_image_path):
    i = 0
    for image in images:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # cv2.imwrite(gray_image_path + 'gray_' + images_name[i].split('\\')[-1], image)
        cv2.imwrite(gray_image_path + 'gray_' + images_name[i].split('/')[-1], image)
        i = i + 1


def threshold_adaptive(image):
    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 中值滤波
    img = cv2.medianBlur(img, 5)

    ret, th1 = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
    # 11 为 Block size, 2 为 C 值
    th2 = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 11, 2)
    th3 = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)

    # titles = ['Original Image', 'Global Threshold (v = 127)', 'Adaptive Mean Threshold', 'Adaptive Gaussian Threshold']
    images = [img, th1, th2, th3]
    return images

def get_binary_images(images, images_name, binary_image_path):
    i = 0
    for image in images:
        images_b = threshold_adaptive(image)
        titles = ['OriginalImage', 'GlobalThreshold', 'AdaptiveMeanThreshold',
                  'AdaptiveGaussianThreshold']
        for image_b in images_b:
            # cv2.imwrite(binary_image_path + 'binary_' + title + '_' + images_name[i].split('\\')[-1], image_b)
            for title in titles:
                cv2.imwrite(binary_image_path + 'binary_' + images_name[i].split('/')[-1].split('.')[0] + '_' + title + '.jpg', image_b)
                # print(binary_image_path + 'binary_' + images_name[i].split('\\')[-1].split('.')[0] + '_' + title + '.jpg')
        i = i + 1


def edge_image(image):
    blurred = cv2.GaussianBlur(image, (3, 3), 0)
    gray = cv2.cvtColor(blurred, cv2.COLOR_BGR2GRAY)

    grad_x = cv2.Sobel(gray, cv2.CV_16SC1, 1, 0)
    grad_y = cv2.Sobel(gray, cv2.CV_16SC1, 0, 1)

    # edge_output = cv2.Canny(grad_x, grad_y, 30, 150)
    edge_output = cv2.Canny(gray, 50, 150)
    return edge_output


def get_edge_images(images, images_name, edge_image_path):
    i = 0
    for image in images:
        image = edge_image(image)
        # print(edge_image_path + 'edge_' + images_name[i].split('\\')[-1])
        cv2.imwrite(edge_image_path + 'edge_' + images_name[i].split('/')[-1], image)
        # cv2.imwrite(edge_image_path + 'edge_' + images_name[i].split('/')[-1], image)
        i = i + 1



if __name__ == "__main__":
    images_path = 'resources/17flowers'
    images, images_name = read_images(images_path)

    # 数据增强
    # 缩放
    # get_resize_images(images, images_name, "./output/resize/")
    # 平移
    # get_trans_images(images, images_name, "./output/trans/")
    # 旋转
    # get_rotate_images(images, images_name, "./output/rotate/")
    # 高斯噪声
    # get_GaussianNoise_images(images,images_name, "./output/gaussian_noise/")
    # 灰度
    # get_gray_images(images, images_name, "./output/gray/")
    # 二值
    # get_binary_images(images, images_name, "./output/binary/")
    # 边缘
    # get_edge_images(images, images_name, "./output/edge/")


    # cv2.waitKey(0)  # 等有键输入或者1000ms后自动将窗口消除，0表示只用键输入结束窗口
    # cv2.destroyAllWindows()  # 关闭所有窗口
